package bm.com.ai.server.controller;

import bm.com.ai.server.dto.EmbeddingStoreDTO;
import bm.com.ai.server.vector.config.EmbeddingModelConfig;
import bm.com.ai.server.vector.service.QdrantVectorService;
import bm.com.framework.common.config.EnumRoute;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import jakarta.annotation.Resource;
import org.springframework.web.bind.annotation.*;

import java.util.List;

@RestController
@RequestMapping("/common" + EnumRoute.AIRoute.Base)
public class VectorModelController {

    @Resource
    EmbeddingModelConfig embeddingModelConfig;
    
    @Resource
    QdrantVectorService qdrantVectorService;

    @Resource
    QdrantEmbeddingStore embeddingStore;

    @PostMapping("/vector-text")
    public String testVectorStore(@RequestBody EmbeddingStoreDTO dto) {
        try {

            if (dto.getText() == null){
                return "向量存储未初始化！";
            }

            TextSegment segment1 = TextSegment.from(dto.getText());
            segment1.metadata().put("source", "test");
            Embedding embedding1 = embeddingModelConfig.embeddingModel().embed(segment1).content();
            String message = embeddingStore.add(embedding1, segment1);

            System.out.println(message);
            return "向量存储测试成功！添加的文档ID: " + message;
        } catch (Exception e) {
            return "向量存储测试失败: " + e.getMessage();
        }
    }

    @PostMapping("/test-search-vector")
    public String testSearchVector(@RequestBody EmbeddingStoreDTO dto) {
        try {

//            String prompt = "鹅 鹅 鹅，曲项向天歌，白毛浮绿水，红掌拨浅薄";
            String prompt = dto.getText();
            List<EmbeddingMatch<TextSegment>> matches = qdrantVectorService.getEmbeddingMatch(prompt);

            System.out.println(matches.size());
            System.out.println("metadata :" + matches.get(0).embedded().metadata());
            return matches.get(0).embedded().text();
        } catch (Exception e) {
            return "向量存储测试失败: " + e.getMessage();
        }
    }

}
